# pre packages
from myglobal import os, sys as sys_global, LINES

# sys packages
from pylab import np
import obspy as ob
from glob import glob
from tqdm import tqdm
from obspy.core.utcdatetime import UTCDateTime
import pandas as pd
from scipy.stats import linregress
import datetime
from scipy import interpolate,signal

# self packages
from utils.loc import load_loc, get_distance,sort_data_by_distance
from utils.math import measure_shift_fft,remove_point_skip, analyze_frequency_components, norm
from utils.h5data import get_event_data, save_h5_data, read_h5_data
from utils.plot import plot_scatters,plot_traces
from utils.trace import get_tapered_slices, safe_filter,get_tapered_traces

def get_freq_by_interpolate(t,x,dt, fs,fe):

    t1 = np.arange(t[0],t[-1],dt)
    f = interpolate.interp1d(t,x, kind='linear', bounds_error=False, fill_value=0)
    y1 = f(t1)
    y1 = signal.detrend(y1)
    y1 = get_tapered_traces(y1,dt,L_Taper=30)

    N = len(y1)
    freqs = np.fft.fftfreq(N, dt)
    P = np.fft.fft(y1)

    fsp = np.argmin(np.abs(freqs-fs))
    fep = np.argmin(np.abs(freqs-fe))
    freqs = freqs[fsp:fep]
    P = P[fsp:fep]
    return freqs,P/N,np.abs(P), np.angle(P)

def main():
    # 解析命令行参数
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-input', default='', help='输入的文件，从8.py来')
    parser.add_argument('-key', default='key',  help='对哪个文件进行分析')
    parser.add_argument('-emarker', default='', help='marker for file')
    parser.add_argument('-figroot', default='figures/9.freq.analysis.figures',  help='root to save figs')
    args = parser.parse_args()
    print(args)
    
    EMARKER = args.emarker
    INPUT = args.input
    FIG_ROOT = args.figroot
    key = args.key

    datasets,args_infile= read_h5_data(INPUT,['all_groups','te'], group_name='metadata', read_attrs=True)
    STATS = [i.decode() for i in datasets[0]]
    te = datasets[1]
    dte = te[1]-te[0]
    print(args_infile)
    DATE=args_infile['date']
    EMARKER = args_infile['emarker']
    ns = len(STATS)

    FIG_ROOT = f'{FIG_ROOT}/9.{DATE}.{EMARKER}'
    if not os.path.exists(FIG_ROOT):
        os.makedirs(FIG_ROOT)

    if DATE == '2303':
        s_info = load_loc('./loc/loc_all_from_log_2303.csv')
        START_TIME = '2023-03-10T00:00:00Z'
    elif DATE == '2406':
        pass
    elif DATE == '2409':
        pass
    elif DATE == '254C':
        s_info = load_loc('./loc/254C.csv',lat_key='lat',lon_key='lon')
        START_TIME = '2024-10-08T00:00:00Z'
    else:
        pass
    
    print('use', INPUT)
    
    P_all = []
    x = np.zeros([ns])
    dt_all = {}

    for j, name in enumerate(STATS):
        xi, azi,_ = get_distance(s_info=s_info,name1=SOURCE,name2=name, S2N=False)
        # 提取当前台站的数据
        
        dt_j = remove_point_skip(delta_t_shift[:, j], peroid=100)
        tmax_j = T_MAX[j]
        dt_j[np.abs(dt_j)>70]=0
        weights = np.ones(len(dt_j))
        weights[dt_j==0]=0
        # 分析频率成分
        # dt_j = np.abs(signal.hilbert(dt_j))
        # dt_j = get_tapered_traces(dt_j,1,L_Taper=1000)
        dt_j = signal.detrend(dt_j)

        freqs, P, A, phi = analyze_frequency_components(
             te_bjt,dt_j.copy()/(0.1+tmax_j), freqs=freqs, weights=weights  # 分析0.1-4次/天的频率范围
        )

        # freqs, P, A, phi = get_freq_by_interpolate(te_bjt,dt_j.copy(),
        #                             dt=1/24,fs=1/50,fe=4)


        P_all.append(P)
        x[j]= xi

        dt_all[name]=dt_j
        
    P_all = np.array(P_all)
    from pylab import figure, plt
    from matplotlib import rcParams, gridspec
    rcParams['font.family'] = 'Arial'
    rcParams['font.size'] = '8'
    DPI=600
    SCALE = 200/50  # m per ms
    plt.close('all')
    fig = figure(figsize=[8,6], dpi=DPI)
    # 绘制频率分析图
    gs = fig.add_gridspec(2,4,wspace=0)
    ax1 = fig.add_subplot(gs[0,:])
    
    for j, name in enumerate(stats):
        dt_j = dt_all[name]
        dt_j[dt_j==0]=np.nan
        ax1.scatter(te_bjt, dt_j*SCALE+x[j], s=10, c='gray', marker='^', edgecolors='none')
        ax1.plot(te_bjt, dt_j*SCALE+x[j], '-', color='tab:red', lw=0.5, label=f'{name}')
        ax1.text(0,dt_j[j]*SCALE+x[j], name, fontsize=8)
    
    # dt_phy,phy, P_phy = get_solid(START_TIME, ndays=180,freqs=freqs)

    # dt_phy,phy, P_phy = get_phy_4C(START_TIME, ndays=180,freqs=freqs, file = './loc/P320T.csv')
    dt_phy,phy, P_phy = get_phy_4C(START_TIME, ndays=180,freqs=freqs, file = 'loc/skt.DZ064.csv')
    # dt_phy,phy, P_phy = get_era5_T(START_TIME, ndays=180,freqs=freqs, file = 'loc/era5.T.grib',key='t2m')
    

    # phy=interpolate.interp1d(dt_phy, phy, kind='linear', bounds_error=False,fill_value=0)(te_bjt)
    # dt_phy = te_bjt
    # print(phy.min())
    # print(phy.max())
    # assert 1==2
    ax1.plot(dt_phy, (norm(phy.copy()))*200-200, '-', color='k', lw=0.5, label='SKT')
    # resp = np.loadtxt('loc/resp.txt',delimiter=' ')
    # dt_fit = interpolate.interp1d(resp[:,0], resp[:,1], kind='linear', bounds_error=False,fill_value=0)(phy)
    # # ax1.plot(dt_phy, (norm(phy.copy()))*400+3500, '-', color='k', lw=0.5, label='SKT')
    # ax1.plot(te_bjt, (dt_all['P300']-dt_fit)*SCALE+3500, '-', color='k', lw=0.5, label='remains')
    # ax1.plot(te_bjt, (dt_fit)*SCALE+4000, '-', color='blue', lw=0.5, label='fit')

    
    ax1.set_xlabel('days(BJT)')
    ax1.set_ylabel('x(m)')
    # ax1.set_xticks(np.arange(0,175,30))
    ax1.set_xticks(np.arange(0,20,7))
    ax1.set_ylim([x.min()-400,x.max()+400])
    # ax1.set_xlim([140,160])

    ax2 = fig.add_subplot(gs[1,0:2])
    for j, name in enumerate(stats):
        ax2.plot(freqs, np.abs(P_all[j,:])*10+x[j], '-', color='k', lw=0.7, label=f'{name}')
    print(stats)
    ax2.plot(freqs, np.abs(P_all[6:,:].sum(axis=0))*1, '-', color='tab:red', lw=0.7, label=f'{name}')

    # ax2.plot(freqs, np.abs(P_tide)*400, '-', color='green', lw=0.5, label=f'{name}')
    ax2.plot(freqs, np.abs(P_phy)*100-400, '-', color='green', lw=0.5, label=f'{name}')

    ax2.scatter([1]*len(x),np.abs(P_all[:, np.argmin(np.abs(freqs-1))])*10+x, s=10, c='tab:red', marker='^', edgecolors='none')
    ax2.plot([1-1/14,1-1/14],[x.min(),x.max()],c='tab:blue', lw=0.4)
    ax2.plot([2-1/14,2-1/14],[x.min(),x.max()],c='tab:blue', lw=0.4)
    ax2.plot([1/7,1/7],[x.min(),x.max()],c='tab:blue', lw=0.4)
    ax2.plot([1/14,1/14],[x.min(),x.max()],c='tab:blue', lw=0.4)
    ax2.scatter([2]*len(x),np.abs(P_all[:, np.argmin(np.abs(freqs-2))])*10+x, s=10, c='tab:blue', marker='o', edgecolors='none')

    ax2.set_xlabel('freq(cycle/day)')
    ax2.set_ylabel('x(m)')
    ax2.set_ylim([x.min()-400,x.max()+400])
    # ax2.set_xscale('log')
    

    ax3 = fig.add_subplot(gs[1,2])

    ax3.scatter(np.angle(P_all[:, np.argmin(np.abs(freqs-1))])/np.pi, x, s=10, c='tab:red', marker='^', edgecolors='none')
    ax3.scatter(np.angle(P_all[:, np.argmin(np.abs(freqs-(1-1/14)))])/np.pi, x, s=10, c='gray', marker='^', edgecolors='none')
    ax3.scatter(np.angle(P_all[:, np.argmin(np.abs(freqs-(1+1/14)))])/np.pi, x, s=10, c='k', marker='^', edgecolors='none')
    ax3.scatter(np.angle(P_all[:, np.argmin(np.abs(freqs-2))])/np.pi, x, s=10, c='tab:blue', marker='o', edgecolors='none')

    ax3.scatter([np.angle(P_phy[np.argmin(np.abs(freqs-1))])/np.pi], [3500], s=10, c='tab:red', marker='*', edgecolors='none')
    ax3.scatter([np.angle(P_phy[np.argmin(np.abs(freqs-2))])/np.pi], [3500], s=10, c='tab:blue', marker='*', edgecolors='none')

    ax3.set_xlabel(r'$\phi$($\pi$)')
    ax3.set_ylim([x.min(),x.max()])
    ax3.set_xlim([-1,1])
    ax3.set_yticks([])

    ax4 = fig.add_subplot(gs[1,3])

    ax4.scatter(np.abs(P_all[:, np.argmin(np.abs(freqs-1))]), x, s=10, c='tab:red', marker='^', edgecolors='none')
    ax4.scatter(np.abs(P_all[:, np.argmin(np.abs(freqs-2))]), x, s=10, c='tab:blue', marker='o', edgecolors='none')
    ax4.set_xlabel('A(ms)')
    ax4.set_ylim([x.min(),x.max()])
    ax4.set_xlim([0,10])
    ax4.set_yticks([])
    
    plt.suptitle(f'{EMARKER}')
    # X,Y = np.meshgrid(t, freqs)

    figname = f'{FIG_ROOT}/freq_analysis.{EMARKER}.png'
    print(figname)
    fig.tight_layout()
    fig.savefig(figname)

    plt.close('all')
    fig = figure(figsize=[8,8], dpi=DPI)
    # phy = norm(phy)*20
    for k in range(60,60+36,1):
        
        phy1 = interpolate.interp1d(dt_phy, phy, kind='linear', bounds_error=False,fill_value=0)(te_bjt)
        phy2 = (te_bjt%1)*24
        # phy2 = phy1
        dt = dt_all['DZ069']
        plt.subplot(6,6,k-60+1)
        mask =(te_bjt > k) & (te_bjt < k+1)
        # plt.scatter(phy1, dt,s=10, c='gray', marker='*', edgecolors='none')
        
        plt.plot(phy2[mask], dt[mask], lw=0.3, color='tab:red')
        # plt.plot(phy2[mask], phy1[mask], lw=0.3, color='tab:blue')
        scatter = plt.scatter(phy2[mask], dt[mask], s=15, c=24*(te_bjt[mask]%1), 
                             cmap='coolwarm', marker='^', edgecolors='none')
        # cbar = plt.colorbar(scatter)
        # cbar.set_label('hours in a day')
        # for T in range(-5,30):
        #     C = dt[np.where((phy1>=T) & (phy1<T+1))]

        #     # 拟合一个C的数值分布
        #     if len(C)>20:
        #         den,edges = np.histogram(C, bins=np.linspace(C.min(),C.max(),30))
        #         den = den/den.max()
        #         plt.plot(den+T,(edges[:-1]+edges[1:])/2,lw=0.2, color='tab:red')
        #         plt.fill_betweenx((edges[:-1]+edges[1:])/2,[T]*len(den),T+den/den.max(),alpha=0.2,color='tab:red')
        #     print(T, C.mean())

        # plt.xlabel('T(deg)')
        # plt.ylabel('dt(ms)')

    plt.suptitle(f'T vs dt, delay={k} hours')
    figname = f'{FIG_ROOT}/{k:02d}test.{EMARKER}.png'
    print(figname)
    fig.tight_layout()
    fig.savefig(figname)


if __name__ == '__main__':
    main()
